import torch
from config.hyper import hyper
from predict import Predictor
import os

if __name__ == '__main__':
    data = [
        "But that spasm of irritation by a master intimidator was minor compared with what Bobby Fischer , "
        "the erratic former world chess champion , dished out in March at a news conference in Reykjavik , Iceland .",
        "But Schaap seems as comfortable in that role as Joe Buck , the Fox baseball and football sportscaster who so "
        "clearly benefited from learning beside his father , Jack Buck , the late voice of the St. Louis Cardinals . "
        "''",
        "www.formula1.com August Aug. 1-5 National Corvette Restorers Society Annual Convention , Henry B. Gonzalez "
        "Convention Center , San Antonio .",
        "(248) 269-7672 , www.meadowbrookconcours.org Aug. 6 Formula One Grand Prix , Budapest , Hungary .",
        "www.lemansclassic.com July 9 Italian Car Festival , Stark County Fairgrounds , Canton , Ohio .",
        ]

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    best = "{}/{}".format(os.path.dirname(os.path.abspath(__file__)), "best_model_sequmt/nyt_seq2umt_ops_best")
    predictor = Predictor(best_model=best, device=device)
    predictor.putinData(use_dataset_file=False, sentence_list=data).predict()
    out = predictor.get_result()
    for x in out["predict_triple"]:
        if len(x) > 0:
            print(x)
